#!/usr/bin/env python
# -*- coding: utf-8 -*-

import cv2
import envs
import time
import rospy
import numpy as np

from img_utils import ros_to_cv2, cv2_to_ros

from std_msgs.msg import Float64
from sensor_msgs.msg import Image


class ImgObjectCenter():
    def __init__(self):
        # 图像发布者
        self.img_object_center_pub = rospy.Publisher('/camera/image_raw/img_object_center', Image, queue_size=1)
        # 物体中心点横坐标发布者
        self.object_center_x_pub = rospy.Publisher('/camera/image_raw/object_center_x', Float64, queue_size=1)
        # 物体中心点纵坐标发布者
        self.object_center_y_pub = rospy.Publisher('/camera/image_raw/object_center_y', Float64, queue_size=1)
        # 图像订阅者
        self.image_sub = rospy.Subscriber('/camera/image_raw', Image, callback=self.image_callback)
        # 物体中心点横坐标
        self.object_center_x = Float64()
        # 物体中心点纵坐标
        self.object_center_y = Float64()
        # canny检测相关参数
        self.ratio = 3
        self.canny_threshold = 30
        # 膨胀腐蚀相关参数
        self.kernel_size = 11
        self.iterations = 1
        # 颜色(BGR)
        self.color = (211, 211, 118)

        time.sleep(0.2)
        rospy.spin()

    def handle_object_center(self, img_raw, img_closing):
        """
        物体中心点
        :param img_raw: 图像原图或者压缩后的图像
        :param img_closing: 图像膨胀腐蚀之后的图像
        """
        rospy.logdebug('handle object center...')

        # 边缘点
        contours, _ = cv2.findContours(img_closing, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
        rospy.logdebug('len(contours): %s', len(contours))
        # 从大到小重新排列
        contours_sorted = sorted(contours, key=lambda contour: len(contour), reverse=True)
        cnt = contours_sorted[0]

        # 画出轮廓点
        cv2.drawContours(img_raw, contours_sorted, -1, self.color, 2)

        # 凸包
        hull = cv2.convexHull(cnt)
        cv2.drawContours(img_raw, [hull], -1, self.color, 2)

        # 最小外接矩形框
        rect = cv2.minAreaRect(cnt)
        box = cv2.boxPoints(rect)
        box = np.int0(box)
        cv2.drawContours(img_raw, [box], 0, self.color, 2)
        # 中心点x,y坐标
        center = rect[0]
        center_x = center[0]
        center_y = center[1]
        # 发布物体中心点
        self.object_center_x.data = center_x
        self.object_center_y.data = center_y
        self.object_center_x_pub.publish(self.object_center_x)
        self.object_center_y_pub.publish(self.object_center_y)

        cv2.circle(img_raw, (int(center_x), int(center_y)), 1, self.color, 5)

        self.img_object_center_pub.publish(cv2_to_ros(img_raw, encoding='bgr8'))

        rospy.logdebug('handle object center OK')

    def image_callback(self, data):
        rospy.logdebug('image callback...')

        img = ros_to_cv2(data, encoding='bgr8')
        # canny边缘检测
        canny = cv2.Canny(img, self.canny_threshold, self.canny_threshold * self.ratio)
        # 膨胀腐蚀
        kernel = cv2.getStructuringElement(2, (self.kernel_size, self.kernel_size))
        closing = cv2.morphologyEx(canny, cv2.MORPH_CLOSE, kernel, iterations=self.iterations)
        self.handle_object_center(img, closing)


if __name__ == '__main__':
    rospy.init_node('npu_robot_img_object_center', anonymous=True, log_level=envs.log_level)
    img_object_center = ImgObjectCenter()
